// The MPC formulation for r^2 as in Ni Dong's dissertation.
#include "MPC_EW2D_1.h"
using namespace Ipopt;
// -----------------------------------------------------------------------
/* Constructor. */
MPC_Formulation::MPC_Formulation()
{
	int m,n;
	h_set = 100.0;
	r2_set = 2.25;
	t_dep = 200.0;
	t_k   = 0.0;
	dt = 5.0;
	mode = 20;
	k_max = 100;
	q_h = 1/(h_set*h_set);
	q_r = 1/(r2_set*r2_set);
	
	T_old = 610;
	W_old = 1.0;

	z = new double[mode*mode];
	A1 = new double*[mode];
	A2 = new double*[mode];
	for(m=0;m<mode;m++)
	{
		A1[m] = new double[mode];
		A2[m] = new double[mode];
		for(n=0;n<mode;n++)
		{
			A1[m][n] = 0.0;
			A2[m][n] = 0.0;
		}
	}

	Km2 = new double*[mode];
	double temp = dx*dx*pi*pi;
	double dx = pi/k_max;
	for(m=0;m<mode;m++)
	{
		Km2[m] = new double[mode];
		for(n=0;n<mode;n++)
		{
			Km2[m][n] = (2-2*cos(2*m*dx))/temp;
		}
	}
}

MPC_Formulation::~MPC_Formulation()
{
	delete []z;
	int m;
	for(m=0;m<mode;m++){
		delete [] A1[m];
		delete [] A2[m];
		delete [] Km2[m];
	}
	delete [] A1;
	delete [] A2;
	delete [] Km2;
}

void MPC_Formulation::parameterize(double time, const double* z_cov)
{
	t_k = time;
	C_k = eq_C(T_old,W_old);
	C2_k = eq_C2(T_old,W_old);
	sigma2_k = eq_sigma2(T_old,W_old);
	int m,n;
	double Z_cov_mn;
	for(m=0;m<mode;m++){
		for(n=0;n<mode;n++){
			Z_cov_mn = z_cov[index2D(m,n)];
			z[index2D(m,n)] = Z_cov_mn;
			if(m!=0||n!=0){
				A1[m][n] = Z_cov_mn*eq_K1(C2_k,m,n,dt);
				A2[m][n] = eq_K2(C2_k,m,n,dt)*sigma2_k;
			}
		}
	}

	dt2 = t_dep-t_k-dt;
}
// -----------------------------------------------------------------------

// -----------------------------------------------------------------------
bool MPC_Formulation::get_nlp_info(Index& n, Index& m, Index& nnz_jac_g,
                         Index& nnz_h_lag, IndexStyleEnum& index_style)
{
  // The problem described in MPC_Formulation.hpp has 2 variables, x1, & x2,
  n = 2;

  // one inequality constraint,
  m = 1;

  // 2 nonzeros in the jacobian (one for x1, and one for x2),
  nnz_jac_g = 2;

  // and 2 nonzeros in the hessian of the lagrangian
  // (one in the hessian of the objective for x2,
  //  and one in the hessian of the constraints for x1)
  nnz_h_lag = 4;

  // We use the standard fortran index style for row/col entries
  index_style = C_STYLE;

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::get_bounds_info(Index n, Number* x_l, Number* x_u,
                            Index m, Number* g_l, Number* g_u)
{
  // here, the n and m we gave IPOPT in get_nlp_info are passed back to us.
  // If desired, we could assert to make sure they are what we think they are.
  assert(n == 2);
  assert(m == 1);

  // x1 has a lower bound of -1 and an upper bound of 1
  x_l[0] = 400;
  x_u[0] = 900;

  // x2 has no upper or lower bound, so we set them to
  // a large negative and a large positive number.
  // The value that is interpretted as -/+infinity can be
  // set in the options, but it defaults to -/+1e19
  x_l[1] = 0.1;
  x_u[1] = 2.0;

  // we have one equality constraint, so we set the bounds on this constraint
  // to be equal (and zero).
  g_l[0] = 0.1*h_set/t_dep;
  g_u[0] = 2e19;

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::get_starting_point(Index n, bool init_x, Number* x,
                               bool init_z, Number* z_L, Number* z_U,
                               Index m, bool init_lambda,
                               Number* lambda)
{
  // Here, we assume we only have starting values for x, if you code
  // your own NLP, you can provide starting values for the others if
  // you wish.
  assert(init_x == true);
  assert(init_z == false);
  assert(init_lambda == false);

  // we initialize x in bounds, in the upper right quadrant
  // x[0] = T_old;
  // x[1] = W_old;
  x[0] = 610;
  x[1] = 0.5;

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_f(Index n, const Number* x, bool new_x, Number& obj_value)
{
  // return the value of the objective function
  // J = q_h*(h_set-h_f)^2+q_r*(r2_s-r2_f)^2

  Number T = x[0];
  Number W = x[1];

  double h_f;
  double r2_f;
  
  double C;
  double C2;
  double sigma2;

  C = eq_C(T,W);
  C2 = eq_C2(T,W);
  sigma2 = eq_sigma2(T,W);

  h_f = eq_hf(C);
  r2_f = eq_r2_f(C2,sigma2);
  
  obj_value = q_h*(h_set-h_f)*(h_set-h_f)+q_r*(r2_set-r2_f)*(r2_set-r2_f);

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_grad_f(Index n, const Number* x, bool new_x, Number* grad_f)
{
  // return the gradient of the objective function grad_{x} f(x)

  double T = x[0];
  double W = x[1];

  double C = eq_C(T,W);
  double C2 = eq_C2(T,W);
  double sigma2 = eq_sigma2(T,W);
  double hf = eq_hf(C);
  double r2_f = eq_r2_f(C2,sigma2);
  double dhfdC = eq_dhfdC(C);
  double dCdT  = eq_dCdT(T,W);
  double dr2_fdC2 = eq_dr2_fdC2(C2,sigma2);
  double dC2dT = eq_dC2dT(T,W);
  double dr2_fdsigma2 = eq_dr2_fdsigma2(C2,sigma2);
  double dsigma2dT = eq_dsigma2dT(T,W);
  double dCdW = eq_dCdW(T,W);
  double dC2dW = eq_dC2dW(T,W);
  double dsigma2dW = eq_dsigma2dW(T,W);

  // grad_{x1} f(x): x1 is not in the objective
  grad_f[0] = -2*q_h*(h_set-hf)*dhfdC*dCdT-2*q_r*(r2_set-r2_f)*(dr2_fdC2*dC2dT+dr2_fdsigma2*dsigma2dT);

  // grad_{x2} f(x):
  grad_f[1] = -2*q_h*(h_set-hf)*dhfdC*dCdW-2*q_r*(r2_set-r2_f)*(dr2_fdC2*dC2dW+dr2_fdsigma2*dsigma2dW);

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_g(Index n, const Number* x, bool new_x, Index m, Number* g)
{
  // return the value of the constraints: g(x)
  Number T = x[0];
  Number W = x[1];

  g[0] = eq_C(T,W);

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_jac_g(Index n, const Number* x, bool new_x,
                       Index m, Index nele_jac, Index* iRow, Index *jCol,
                       Number* values)
{
  if (values == NULL) {
    // return the structure of the jacobian of the constraints

    // element at 1,1: grad_{x1} g_{1}(x)
    iRow[0] = 0;    jCol[0] = 0;

    // element at 1,2: grad_{x2} g_{1}(x)
    iRow[1] = 0;    jCol[1] = 1;
  }
  else {
    // return the values of the jacobian of the constraints
	double T = x[0];
	double W = x[1];

    // element at 1,1: grad_{x1} g_{1}(x)
    values[0] = eq_dCdT(T,W);

    // element at 1,2: grad_{x1} g_{1}(x)
    values[1] = eq_dCdW(T,W);
  }

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_h(Index n, const Number* x, bool new_x,
                   Number obj_factor, Index m, const Number* lambda,
                   bool new_lambda, Index nele_hess, Index* iRow,
                   Index* jCol, Number* values)
{
  if (values == NULL) {
    // return the structure. This is a symmetric matrix, fill the lower left
    // triangle only.
    iRow[0] = 0;    jCol[0] = 0;
    iRow[1] = 1;    jCol[1] = 0;
	iRow[2] = 1;    jCol[2] = 1;
  }
  else {
    // return the values
    // fill the objective portion
    values[0] = 0;
    values[1] = 0;
	values[2] = 0;
	// add the portion for the first constraint
	values[0] += 0;
	values[1] += 0;
	values[2] += 0;
  }
  return false; // USE Quasi-Newton method
}
// -----------------------------------------------------------------------
void MPC_Formulation::finalize_solution(SolverReturn status,
                              Index n, const Number* x, const Number* z_L, const Number* z_U,
                              Index m, const Number* g, const Number* lambda,
                              Number obj_value,
			      const IpoptData* ip_data,
			      IpoptCalculatedQuantities* ip_cq)
{
  // here is where we would store the solution to variables, or write to a file, etc
  // so we could use the solution. Since the solution is displayed to the console,
  // we currently do nothing here.

  // Save the current solution, which will be used as the initial solution.

	T_old = x[0];
	W_old = x[1];

	double C;
	double C2;
	double sigma2;

    double h_f;
    double r2_f;

    C = eq_C(T_old,W_old);
    C2 = eq_C2(T_old,W_old);
    sigma2 = eq_sigma2(T_old,W_old);

    h_f = eq_hf(C);
    r2_f = eq_r2_f(C2,sigma2);

	printf("*** t_k = %f, r_f = %f, h_f = %f",t_k,r2_f,h_f);
}



// -----------------------------------------------------------------------
double MPC_Formulation::eq_C(double T,double W)
{
	return W*(1-k_w*pow(W,-a_w)*exp(k_B*T/E_w));
}

double MPC_Formulation::eq_dCdT(double T,double W)
{
	return -W*k_w*exp(k_B*T/E_w)*k_B/(pow(W,a_w)*E_w);
}

double MPC_Formulation::eq_dCdW(double T,double W)
{
	return 1-k_w*exp(k_B*T/E_w)*(1-a_w)*pow(W,-a_w);
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_C2(double T,double W)
{
	return k_c*exp(k_B*T/E_c)*pow(W,-a_c)/(k_max*k_max);
}

double MPC_Formulation::eq_dC2dT(double T,double W)
{
	return k_c/(k_max*k_max*pow(W,a_c))*exp(k_B*T/E_c)*k_B/E_c;
}

double MPC_Formulation::eq_dC2dW(double T,double W)
{
	return k_c/(k_max*k_max)*exp(k_B*T/E_c)*(-a_c)*pow(W,-a_c-1);
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_sigma2(double T,double W)
{
	double index = (a_t+k_t*W)*T-a_v-k_v*W;
	return pi*pi/(k_max*k_max)*(1+exp(index));
}

double MPC_Formulation::eq_dsigma2dT(double T,double W)
{
	return pi*pi/(k_max*k_max)*exp((a_t+k_t*W)*T-a_v-k_v*W)*(a_t+k_t*W);
}

double MPC_Formulation::eq_dsigma2dW(double T,double W)
{
	return 	pi*pi/(k_max*k_max)*exp((a_t+k_t*W)*T-a_v-k_v*W)*(k_t*T-k_v);
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_K1(double C2,int m, int n,double delta_t)
{
	return exp(-8*(m*m+n*n)*C2*delta_t);
}

double MPC_Formulation::eq_dK1dC2(double C2,int m,int n,double delta_t)
{
	return eq_K1(C2,m,n,delta_t)*(-8*(m*m+n*n)*delta_t);
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_K2(double C2,int m,int n,double delta_t)
{
	return (exp(-8*(m*m+n*n)*C2*delta_t)-1)/(-8*(m*m+n*n)*C2);
}

double MPC_Formulation::eq_dK2dC2(double C2,int m,int n,double delta_t)
{
	double K = -8*(m*m+n*n);
	double temp1 = exp(K*delta_t*C2);
	return (temp1*K*delta_t*K*C2-K*(temp1-1))/(K*K*C2*C2);
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_r2_f(double C2,double sigma2)
{
	int m,n;

	double r2_f = 0.0;
	for(m=0;m<mode;m++){
		for(n=0;n<mode;n++){
			if(m!=0||n!=0){
				r2_f += A1[m][n]*eq_K1(C2,m,n,dt2)+A2[m][n]*eq_K1(C2,m,n,dt2)+sigma2*eq_K2(C2,m,n,dt2);
			}
		}
	}
	return r2_f/=(pi*pi);
}

double MPC_Formulation::eq_dr2_fdC2(double C2,double sigma2)
{
	int m,n;
	double dr2_fdC2 = 0.0;
	for(m=0;m<mode;m++){
		for(n=0;n<mode;n++){
			if(m!=0||n!=0){
				dr2_fdC2 += A1[m][n]*eq_dK1dC2(C2,m,n,dt2)+A2[m][n]*eq_dK1dC2(C2,m,n,dt2)+sigma2*eq_dK2dC2(C2,m,n,dt2);
			}
		}
	}
	dr2_fdC2 /= pi*pi;
	return dr2_fdC2;
}

double MPC_Formulation::eq_dr2_fdsigma2(double C2,double sigma2)
{
	int m,n;
	double dr2_fdsigma2 = 0.0;
	for(m=0;m<mode;m++){
		for(n=0;n<mode;n++){
			if(m!=0||n!=0){
				dr2_fdsigma2 += eq_K2(C2,m,n,dt2);
			}
		}
	}
	dr2_fdsigma2 /= pi*pi;
	return dr2_fdsigma2;
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_hf(double C)
{
	return sqrt(z[index2D(0,0)])/pi+C_k*dt+C*dt2;
}

double MPC_Formulation::eq_dhfdC(double C)
{
	return dt2;
}
// -----------------------------------------------------------------------
void MPC_Formulation::get_solution(double *x)
{
	x[0] = T_old;
	x[1] = W_old;
}